package com.xsdl.flyweight.buffer;

import java.nio.ByteBuffer;
import java.util.ArrayDeque;
import java.util.Deque;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

public class BufferPool {

    private final int totalSize;
    private final int slotSize;
    private int free;
    private final Deque<ByteBuffer> freeSlots = new ArrayDeque<>();
    private final Deque<Condition> waiters = new ArrayDeque<>();
    private final Lock lock = new ReentrantLock();

    public BufferPool(int totalSize, int slotSize) {
        this.totalSize = totalSize;
        this.slotSize = slotSize;
        this.free = totalSize;
    }

    public ByteBuffer allocate(int size, long timeout) throws InterruptedException, TimeoutException {
        if (size > totalSize || size <= 0) {
            throw new IllegalArgumentException("Invalid size: " + size);
        }
        long remainingTime = timeout;
        long start = System.currentTimeMillis();
        lock.lock();
        try {
            if (size == slotSize && !freeSlots.isEmpty()) {
                return freeSlots.pollFirst();
            }
            if (free + freeSlots.size() * slotSize > size) {
                free(size);
                free -= size;
                return ByteBuffer.allocateDirect(size);
            }
            Condition condition = lock.newCondition();
            waiters.addLast(condition);
            try {
                while (true) {
                    boolean wakeUp = condition.await(remainingTime, TimeUnit.MILLISECONDS);
                    if (!wakeUp) {
                        throw new TimeoutException("Timeout waiting for buffer allocation");
                    }
                    if (size == slotSize && !freeSlots.isEmpty()) {
                        return freeSlots.pollFirst();
                    }
                    if (free + freeSlots.size() * slotSize > size) {
                        free(size);
                        return ByteBuffer.allocateDirect(size);
                    }
                    remainingTime = timeout - (System.currentTimeMillis() - start);
                }
            } finally {
                waiters.remove(condition);
            }
        } finally {
            if (!waiters.isEmpty() && !(free == 0 && freeSlots.isEmpty())) {
                waiters.peekFirst().signal();
            }
            lock.unlock();
        }
    }

    private void free(int size) {
        while (free < size) {
            free += freeSlots.pollFirst().capacity();
        }
    }

    public void deallocate(ByteBuffer buffer) {
        lock.lock();
        try {
            if (buffer.capacity() == slotSize) {
                buffer.clear();
                freeSlots.addLast(buffer);
            } else {
                free += buffer.capacity();
            }
            if (!waiters.isEmpty()) {
                waiters.peekFirst().signal();
            }
        } finally {
            lock.unlock();
        }
    }

}
